import os
import uvicorn
from pathlib import Path

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from multiprocess import cpu_count
from contextlib import asynccontextmanager
from urllib.parse import urlencode

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware

from logger import configure
from loguru import logger

class ReqQuestion(BaseModel):
    question: str = Field(default=None, nullable=False, description="Question")

class RespAnswer(BaseModel):
    answer: str = Field(default=None, nullable=False, description="Answer")


def initialize_services():
    pass

def teardown_services():
    pass

@asynccontextmanager
async def lifespan(app: FastAPI):
    initialize_services()
    yield
    teardown_services()

def create_app():
    app = FastAPI(lifespan=lifespan)

    origins = ["*"]

    app.add_middleware(
        CORSMiddleware,
        allow_origins=origins,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    import traceback
    @app.middleware("http")
    async def catch_exceptions_middleware(request: Request, call_next):
        try:
            return await call_next(request)
        except Exception as e:
            logger.error(traceback.format_exc())
            raise e
    
    @app.middleware("http")
    async def flatten_query_string_lists(request: Request, call_next):
        flattened: list[tuple[str, str]] = []
        for key, value in request.query_params.multi_items():
            flattened.extend((key, entry) for entry in value.split(","))

        request.scope["query_string"] = urlencode(flattened, doseq=True).encode("utf-8")

        return await call_next(request)

    @app.get("/health")
    def health():
        logger.info("Health check")
        return {"status": "ok"}

    @app.post("/query", response_model=RespAnswer)
    def query(q: ReqQuestion):
        from search.kg import ChronRAG
        k = ChronRAG()
        return RespAnswer(answer=f"{k.invoke(q.question)}")

    return app

def get_number_of_workers(workers=None):
    if workers == -1 or workers is None:
        workers = (cpu_count() * 2) + 1
    logger.debug(f"Number of workers: {workers}")
    return workers

if __name__ == "__main__":
    # Load environment variables
    load_dotenv()
    
    # Get configuration from environment variables
    HOST = os.getenv("HOST", "0.0.0.0")
    PORT = int(os.getenv("PORT", "8000"))
    APP_NAME = os.getenv("APP_NAME", "ChronRAG")
    LOG_LEVEL = os.getenv("LOG_LEVEL", "debug")
    LOG_FILE = os.getenv("LOG_FILE", "logs/app.log")
    WORKERS = int(os.getenv("WORKERS", get_number_of_workers()))
    
    configure(log_level=LOG_LEVEL, log_file=Path(LOG_FILE))

    app = "main:create_app"
    
    logger.debug(f"{APP_NAME} starting.... {HOST}:{PORT}")
    uvicorn.run(
        app, 
        host=HOST, 
        port=PORT, 
        workers=WORKERS, 
        log_level=LOG_LEVEL
    )
